/*
 *  Copyright (C) 2020 original author or authors.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 *  in compliance with the License. You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software distributed under the License
 *  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 *  or implied. See the License for the specific language governing permissions and limitations under
 *  the License.
 */

package vip.justlive.easyproxy.client.handler;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.Map;
import java.util.function.BiConsumer;
import lombok.extern.slf4j.Slf4j;
import vip.justlive.easyproxy.core.context.ClientContext;
import vip.justlive.easyproxy.core.context.ClientProxyContext;
import vip.justlive.easyproxy.core.context.Config;
import vip.justlive.easyproxy.core.protocol.AbstractProxyDataHandler;
import vip.justlive.easyproxy.core.protocol.ChannelData;
import vip.justlive.easyproxy.core.protocol.ProxyData;
import vip.justlive.easyproxy.core.protocol.ProxyDataType;
import vip.justlive.oxygen.core.config.ConfigFactory;
import vip.justlive.oxygen.core.exception.Exceptions;
import vip.justlive.oxygen.core.net.aio.core.AioListener;
import vip.justlive.oxygen.core.net.aio.core.ChannelContext;
import vip.justlive.oxygen.core.net.aio.core.Client;
import vip.justlive.oxygen.core.net.aio.core.GroupContext;
import vip.justlive.oxygen.core.util.SystemUtils;

/**
 * 客户端处理单元
 *
 * @author wubo
 */
@Slf4j
public class ClientHandler extends AbstractProxyDataHandler implements AioListener {

  private final EnumMap<ProxyDataType, BiConsumer<ProxyData, ChannelContext>> consumers = new EnumMap<>(
      ProxyDataType.class);
  private final Map<Long, Long> relation = new HashMap<>();
  private ClientContext context;
  private final Config config;

  public ClientHandler(Config config) {
    this.config = config;
    consumers.put(ProxyDataType.CONNECT, this::handleConnect);
    consumers.put(ProxyDataType.TRANSFER, this::handleTransfer);
    consumers.put(ProxyDataType.CLOSE_CHANNEL, this::handleCloseChannel);
    consumers.put(ProxyDataType.CLOSE_PROXY, this::handleCloseProxy);
    consumers.put(ProxyDataType.CLOSE, this::handleClose);
  }

  @Override
  public void doHandle(ProxyData data, ChannelContext channelContext) {
    BiConsumer<ProxyData, ChannelContext> consumer = consumers.get(data.getType());
    if (consumer == null) {
      log.warn("no consumer found for {}", data.getType());
      return;
    }
    consumer.accept(data, channelContext);
  }

  @Override
  public void onConnected(ChannelContext channelContext) {
    context = new ClientContext().setChannel(channelContext);
    channelContext.write(ProxyData.auth(config.getServer().getAuthKey()));
  }

  public void stats() {
    if (context == null) {
      return;
    }

    for (ClientProxyContext proxy : context.getProxies()) {
      int size = proxy.getClient().getGroupContext().getChannels().size();
      if (size > 0) {
        log.info("proxy-[{}/{}], clients-[{}]", proxy.getId(), proxy.getAddress(), size);
      }
    }
  }

  private void handleConnect(ProxyData data, ChannelContext channelContext) {
    ChannelData cd = data.deserialize();
    if (cd.getData() == null || cd.getData().length == 0) {
      log.error("address is illegal {}", channelContext);
      return;
    }
    String address = new String(cd.getData(), StandardCharsets.UTF_8);
    if (context == null) {
      log.error("client is not init {}", channelContext);
      return;
    }
    int id = data.getId();
    ClientProxyContext proxy = context.find(id);
    if (proxy == null) {
      proxy = new ClientProxyContext().setId(id).setAddress(address);
      context.getProxies().add(proxy);
    }
    getClientProxy(proxy, cd.getId());
  }

  private void handleTransfer(ProxyData data, ChannelContext channelContext) {
    ClientProxyContext proxy = context.find(data.getId());
    if (proxy == null) {
      log.warn("transfer: proxy not found {} {}", data.getId(), channelContext);
      return;
    }

    ChannelData cd = data.deserialize();
    getClientProxy(proxy, cd.getId()).write(cd.getData());
  }

  private void handleCloseChannel(ProxyData data, ChannelContext channelContext) {
    ClientProxyContext proxy = context.find(data.getId());
    if (proxy == null) {
      log.warn("close proxy channel: proxy not found {}, {}", data.getId(), channelContext);
      return;
    }

    ChannelData cd = data.deserialize();
    Long target = relation.remove(cd.getId());
    if (target == null) {
      return;
    }
    ChannelContext client = proxy.getClient().getGroupContext().getChannels().get(target);
    if (client != null) {
      log.info("close proxy channel: {}/{} {}", proxy.getId(), proxy.getAddress(), client);
      client.close();
    }
  }

  private void handleCloseProxy(ProxyData data, ChannelContext channelContext) {
    ClientProxyContext proxy = context.remove(data.getId());
    if (proxy == null) {
      log.warn("close proxy: proxy not found {}, {}", data.getId(), channelContext);
      return;
    }

    for (ChannelContext channel : proxy.getClient().getGroupContext().getChannels().values()) {
      Long id = (Long) channel.getAttr(Config.REMOTE);
      relation.remove(id);
    }
    proxy.getClient().close();
    proxy.setClient(null);
    log.info("close proxy: {}", proxy);
  }

  private void handleClose(ProxyData proxyData, ChannelContext channelContext) {
    log.info("close client: {}", channelContext);
    channelContext.close();
    channelContext.getGroupContext().close();
  }

  private synchronized ChannelContext getClientProxy(ClientProxyContext proxy, long id) {
    Long target = relation.get(id);
    if (proxy.getClient() == null) {
      ProxyHandler handler = new ProxyHandler(proxy.getId(), context.getChannel());
      GroupContext ctx = new GroupContext(handler).setAioListener(handler);
      ConfigFactory.load(ctx, Config.CLIENT_PROXY_AIO_CONFIG);
      proxy.setClient(new Client(ctx));
    }
    ChannelContext ctx;
    if (target == null) {
      try {
        ctx = proxy.getClient().connect(SystemUtils.parseAddress(proxy.getAddress()));
        ctx.join();
      } catch (IOException e) {
        log.error("", e);
        throw Exceptions.wrap(e);
      }
      ctx.addAttr(Config.REMOTE, id);
      relation.put(id, ctx.getId());
    } else {
      ctx = proxy.getClient().getGroupContext().getChannels().get(target);
    }
    if (ctx == null) {
      throw Exceptions.fail("not found proxy channel");
    }
    return ctx;
  }
}
